import io
from PIL import Image
from detectron2.engine import DefaultPredictor
from detectron2.model_zoo import model_zoo
from detectron2.data import MetadataCatalog
from detectron2.config import get_cfg
from detectron2.structures import Boxes, BoxMode
import logging
import os
import json
import torch
import numpy as np

logger = logging.getLogger(__name__)


class ObjectPredictor:
    """
    ObjectDetector handler class. This handler takes an image
    and returns list of detected classes and bounding boxes respectively
    """

    def __init__(self):
        self.mapping = None
        # only cpu
        self.predictor = None
        self.device = torch.device("cpu")
        self.initialized = False
        self.manifest = None

    def initialize(self, ctx):

        self.initialized = False
        self.manifest = ctx.manifest
        properties = ctx.system_properties
        model_dir = properties.get("model_dir")
        # Read the mapping file, index to object name
        mapping_file_path = os.path.join(model_dir, "index_to_name.json")

        if os.path.isfile(mapping_file_path):
            with open(mapping_file_path) as f:
                self.mapping = json.load(f)
        else:
            logger.warning('Missing the index_to_name.json file. Inference output will not include class name.')
        CONFIG_PATH = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
        self.cfg = get_cfg()
        self.cfg.merge_from_file(model_zoo.get_config_file(CONFIG_PATH))
        serialized_file = self.manifest['model']['serializedFile']
        model_pt_path = os.path.join(model_dir, serialized_file)
        self.cfg.MODEL.WEIGHTS = model_pt_path
        self.cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
        self.cfg.MODEL.DEVICE = "cpu"
        labels = self.mapping['object_type_names']
        self.cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(labels)
        MetadataCatalog.get("coco_test").set(
            thing_classes=labels
        )
        self.predictor = DefaultPredictor(self.cfg)
        self.initialized = True

    def preprocess(self, data):
        """
         Scales, crops, and normalizes a image for a PyTorch model,
         returns an Numpy array
        """
        image = data[0].get("data")
        print(image)
        if image is None:
            image = data[0].get("body")
        image = Image.open(io.BytesIO(image))
        return np.asarray(image, np.uint8)

    def inference(self, img):
        predictions = self.predictor(img)
        return predictions["instances"]

    def postprocess(self, inference_output):

        boxes = inference_output.pred_boxes.tensor.numpy()
        boxes = BoxMode.convert(boxes, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
        boxes = boxes.tolist()
        scores = inference_output.scores.tolist()
        pred_classes = inference_output.pred_classes.tolist()
        retval = []
        try:
            if self.mapping:
                labels = [self.mapping['object_type_names'][i] for i in pred_classes]  # Get the Prediction Score
                return retval.append({"labels": labels, "boxes": boxes, "scores": scores})
        except Exception as e:
            raise Exception('Object name list file should be json format - {"object_type_names":["person","car"...]}"')


_detector = ObjectPredictor()


def handle(img, context):
    try:
        if not _detector.initialized:
            _detector.initialize(context)

        if img is None:
            return None

        img = _detector.preprocess(img)
        logger.info(img)
        predictions = _detector.inference(img)
        logger.info(predictions)
        return _detector.postprocess(predictions)

    except Exception as e:
        raise Exception("Please provide a custom handler in the model archive.")
